import torch
import torch.nn as nn
from model.conv_bn_relu import ConvBNRelu

'''
encoder: RBG-image
'''

class Encoder(nn.Module):
    """
    Inserts a watermark into an image.
    """
    def __init__(self, channels, H, W, M):
        super(Encoder, self).__init__()
        self.channels = channels

        self.H = H
        self.W = W
        self.M = M
        self.layers1 = nn.Sequential(ConvBNRelu(3, self.channels))
        self.layers2 = nn.Sequential(ConvBNRelu(self.channels + self.M, self.channels))
        self.layers3 = nn.Sequential(ConvBNRelu(2 * self.channels + self.M, self.channels))
        self.layers4 = nn.Conv2d(3 * self.channels + self.M, 3, kernel_size=1)

    def forward(self, image, message):

        # First, add two dummy dimensions in the end of the message.
        # This is required for the .expand to work correctly
        expanded_message = message.unsqueeze(-1)
        expanded_message.unsqueeze_(-1)

        expanded_message = expanded_message.expand(-1,-1, self.H, self.W) # (batchsize, message_length, H, W)~(64, 30, 128, 128)

        encoded_image1 = self.layers1(image) # (64, 64, 128, 128)
        concat1 = torch.cat([expanded_message, encoded_image1], dim=1) # (64,64+30,128,128)

        encoded_image2 = self.layers2(concat1) # (64, 64, 128, 128)
        concat2 = torch.cat([expanded_message, encoded_image1, encoded_image2], dim=1) #(64,64*2+30,128,128)

        encoded_image3 = self.layers3(concat2) # (64, 64, 128, 128)
        concat3 = torch.cat([expanded_message, encoded_image1, encoded_image2, encoded_image3], dim=1) #(64,64*3+30,128,128)

        encoded_image4 = self.layers4(concat3) #(64,3,128,128)
        encoded_image = encoded_image4 + image

        return encoded_image